rm(list = ls())
library(pacman)
p_load(numDeriv, devtools, matlab, Matrix, nleqslv, ggplot2, control)
load_all()

# Parameterization
delta <- 0.08
alpha <- 0.36
n <- 0
beta0 <- 0.96
sigma <- 2.0
# sigma=1.0
periods <- 40	# number of transition periods

# steady state values
kss <- (alpha/(1/beta0-1+delta))^(1/(1-alpha))
yss <- production(1,kss)
css <- yss-(n+delta)*kss

# stability of steady state
print("stability in the system with k_t+1, k_t, k_t-1")
print("Jac: ")
Jac <- jacobian(diff_heer, x = c(kss,kss))
Jac
eigen(Jac)$values

print("stability in the system with k_t+1, c_t+1")
print("Jac: ")
Jac <- jacobian(diff1_heer, x = c(kss,css))
print("eigenvalues: ")
eigen(Jac)$values

print("analytical Jacobian")
jac1 <- zeros(2,2)
jac1[1,1] <- 1/(1+n)*(alpha*kss^(alpha-1)+1-delta)
jac1[1,2] <- -1/(1+n)
jac1[2,1] <- css*beta0*1/sigma*alpha*(alpha-1)/(1+n)*kss^(alpha-2)*(alpha*kss^(alpha-1)+1-delta)
jac1[2,2] <- 1-css*beta0/sigma*alpha*(alpha-1)/(1+n)*kss^(alpha-2)
print("jac1:")
jac1
eigen(jac1)$values

print("schur factorization: ")
ans <- Schur(jac1)
ans$T
# 要重新排序，特征值从小到大依次排列
ans <- ordschur(ans$Q, ans$T, order(ans$EValues))
TT <- ans$U
S <- ans$S


kt <- kss*ones(periods,1)
zt <- ones(periods,1)
zt[10:12] <- 1.1*ones(3,1)
dynamics1(kt)
x1 <- nleqslv(kt, dynamics1)$x

consumption1 <- prod1 <- zeros(periods,1)

i <- 0
while(i < periods){
  i <- i+1
  if (i==1){
    prod1[i] <- production(1,kss)
    consumption1[i] <- production(1,kss)+(1-delta)*kss-(1+n)*x1[i]
  } else {
    prod1[i] <- production(zt[i-1],x1[i-1])
    consumption1[i] <- production(zt[i-1],x1[i-1])+(1-delta)*x1[i-1]-(1+n)*x1[i]
  }
}

# dynamics if shock is not known in advance
# 比如在第10期开始冲击，那么前10期都要处于稳态，从11期以后开始解模型
kt1 <- kss*ones(periods-10,1) # 注意只有30期
dynamics2(kt1)
x2 <- nleqslv(kt1, dynamics2)$x
x2 <- c(kss*ones(10,1), x2)

consumption2 <- prod2 <- zeros(periods,1)
i <- 0
while(i < periods){
  i <- i+1
  if (i==1){
    prod2[i] <- production(1,kss)
    consumption2[i] <- production(1,kss)+(1-delta)*kss-(1+n)*x2[i]
  } else {
    prod2[i]=production(zt[i-1],x2[i-1])
    consumption2[i]=production(zt[i-1],x2[i-1])+(1-delta)*x2[i-1]-(1+n)*x2[i]
  }
}

picdata <- data.frame(tt = 1:length(x1), kt = x1, ktu = x2, yt = prod1, ytu = prod2,
                      ct = consumption1, ctu = consumption2)
# Fig 2.1
ggplot(picdata, aes(x = tt, y = kt)) + geom_line() +  geom_line(aes(y = ktu), linetype = 2) +
  labs(x = 'Period t', y = 'Capital stock kt') + theme_bw()
# Fig 2.2
ggplot(picdata, aes(x = tt, y = yt)) + geom_line() +  geom_line(aes(y = ytu), linetype = 2) +
  labs(x = 'Period t', y = 'Production yt') + theme_bw()
# Fig 2.3
ggplot(picdata, aes(x = tt, y = ct)) + geom_line() +  geom_line(aes(y = ctu), linetype = 2) +
  labs(x = 'Period t', y = 'Consumption ct') + theme_bw()
# Fig 2.7
ggplot(picdata[13:40,], aes(x = ktu, y = ctu)) + geom_line()
  labs(x = 'Period t', y = 'Consumption ct') + theme_bw()

# linearization
# ct - c = -t21/t22 * (kt-k)
t21dt22 <- -solve(TT)[2,1]/solve(TT)[2,2]
clin1 <- css+(x2[13:40]-kss)*t21dt22

# simulation of linear adjustment
ktlin <- zeros(40-13+1,1)
n1 <- nrow(ktlin)
ktlin[1] <- x2[13]
ctlin <- ktlin
ctlin[1] <- css+(ktlin[1]-kss)*t21dt22
i <- 1
while(i < n1){
  i <- i+1
  ktlin[i] <- kss+S[1,1]*(ktlin[i-1]-kss)
  ctlin[i] <- css+(ktlin[i]-kss)*t21dt22
}

# Fig 2.8,仅模拟40期
picdata <- data.frame(tt = 1:nrow(ktlin), kt = x2[13:40], ktlin = ktlin)
ggplot(picdata, aes(x = tt, y = kt)) + geom_line() +  geom_line(aes(y = ktlin), linetype = 2) +
  labs(x = 'Period t', y = 'kt') + theme_bw()

















